#coding:utf-8

import time
import subprocess
import importlib
import os

from psutil import Process as Process_
#from subprocess import PIPE

from multiprocessing import Process as pro
from FrameWork.Normalization import FuzzData
from Libs.CommunicationLib import Comm
from Libs.ServerLib import getPort,monitor_process

try:
    from xmlrpc.server import SimpleXMLRPCServer
except:
    from SimpleXMLRPCServer import SimpleXMLRPCServer


class ServerNode():
    def __init__(self ,address ):
        # address = 127.0.0.1:1000
        # address_url = 127.0.0.1
        # address_port = 1000
        self.address_url, self.address_port = getPort(address)      
        self._data = {}
        # fun list
        self._rpc_methods_ = ['importPlugin','FuzzRun','FuzzStat']
        self.run_info = {}
        
        #init run info
        self.run_info['MainRun'] = {}
        self.run_info['MainRun']['pid'] = monitor_process("MainRun")
        self.run_info['MainRun']['isrun'] = True
        self.run_info['MainRun']['runtime'] = time.time()
        #
        
        self.import_str = None
        self.load_dict = {}
        self.fuzz = None
        
        #是否已经使用了import？
        self.is_import = None
        
        self._serv = SimpleXMLRPCServer((self.address_url, self.address_port), allow_none=True)
        for name in self._rpc_methods_:
            self._serv.register_function(getattr(self, name))


    # def _start(self):
    #     if self._serv != None:
    #         self._forever()
    #         return
    #
    #     self._serv = SimpleXMLRPCServer((self.address_url,self.address_port), allow_none=True)
    #     for name in self._rpc_methods_:
    #         self._serv.register_function(getattr(self, name))
    #
    #     self._forever()
    

    def _forever(self):
        self._serv.serve_forever()

    #验证
    def _check(self):
        if self.fuzz == None:
            return False
        return True
        
    
    # 解析插件
    # 生成插件
    # 动态导入插件
    def importPlugin(self,plugin_name,load_dict = {"core":"fuzz","system":"boofuzz"}):
        if self.is_import == True:
            return
        
        fdata = FuzzData()
        
        plugin_path = os.getcwd() + "/Plugin/" + plugin_name
        fdata.load_xml_file(plugin_path + ".xml")
        if load_dict['core'] == "fuzz":
            if load_dict['system'] == "boofuzz":
                func = fdata.boofuzz() # 动态的去选择fuzz工具
            else:
                func = []
        else:
            func = []
        
        x = fdata.str_to_tmpe(func)
        if os.path.exists(plugin_path + ".py"):
            os.remove(plugin_path + ".py")
        
        with open(plugin_path + ".py", "w") as f:
            f.write(x)
        
        
        plugin_name = 'Plugin.%s' % plugin_name
        self.import_str = importlib.import_module(plugin_name)
        self.load_dict = load_dict
        self.is_import = True
        # import Plugin.plugin_name


    #需要通过此方法运行程序
    #process_name 为 需要fuzz的应用
    def FuzzRun(self,process_name,cmd = ""):
        self.run_info[process_name] = {}
        self.run_info[process_name]['isrun'] = False
        if cmd == "":       #cmd 启动命令
            pid_ = monitor_process(process_name)
            if pid_ == None:    #如果没启动
                return False
            
            self.run_info[process_name]['p'] = None
            self.run_info[process_name]['cmd'] = cmd
            self.run_info[process_name]['pid'] = pid_
            self.run_info[process_name]['isrun'] = True
            self.run_info[process_name]['runtime'] = 0
            self.run_info[process_name]['poll'] = None
        else:   #如果存在cmd，启动它
            p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=False,shell=True)
            self.run_info[process_name]['p'] = p
            self.run_info[process_name]['cmd'] = cmd
            self.run_info[process_name]['pid'] = int(p.pid)
            self.run_info[process_name]['isrun'] = True
            self.run_info[process_name]['runtime'] = time.time()
            self.run_info[process_name]['poll'] =p.poll()

        #print(self.run_info[process_name]['pid'])
        # 启动fuzz进行测试
        if self.is_import == True:
            start = self.import_str.Fuzz()
            # 多线程运行
            t = pro(target=start.Run, args=(self.load_dict,))
            #start.Run(self.load_dict)
            t.start()
            self.is_import = False
        else:
            return False
    
        return True

    # process_name 应用名
    def FuzzStat(self,process_name):
        process = self.run_info[process_name]
        #p = process['poll']
        pid = process['pid']
        cmd = process['cmd']
        p = Process_(pid)
        #pid_ = monitor_process(cmd)

        self.run_info[process_name]['poll'] = p.status()

        # if pid_ == None and int(pid_) != int(pid):
        #     return False
        
        #记录工作
        comm = Comm()
        if comm.send_log(process_name) == True:
            os.remove("/tmp/" + process_name + ".fuzz")
            return None
        ##
        
        self.run_info[process_name]['exe'] = p.exe()
        self.run_info[process_name]['pwd'] = p.pwd()
        self.run_info[process_name]['uids'] = p.uids()
        self.run_info[process_name]['gids'] = p.gids()
        self.run_info[process_name]['connectios'] = p.connectios()
        self.run_info[process_name]['num_threads'] = p.num_threads()

        #self.run_info[process_name]['runtime'] = False
        
        return self.run_info[process_name]


